{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "87032c94",
   "metadata": {},
   "source": [
    "# What's Pandera?\n",
    "\n",
    "Pandera is an open source framework for precision data testing, built for\n",
    "data scientists and ML engineers.\n",
    "\n",
    "In this notebook, you'll learn how to:\n",
    "\n",
    "> 1. Define Pandera schemas for your dataframe-like objects 📦\n",
    "> 2. Integrate them seamlessly into your data pipelines 🔀\n",
    "> 3. Ensure your data and data transformation functions are correct ✅\n",
    "\n",
    "▶️ Follow the tutorial and run the code cells below to get a sense of how Pandera works\n",
    "and how its error reporting system can provide direct insight into what specific data\n",
    "values caused the error."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5d6cd7a4",
   "metadata": {},
   "source": [
    "First, install pandera:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dabb52d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pandera"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "df888aa0",
   "metadata": {},
   "source": [
    "## What are Schemas?\n",
    "\n",
    "Dataframes and dataframe-like objects are structures with expected properties or rules for the data contained inside. \n",
    "Most of the rules of these structures are known by the designers or analysts of the data, but not directly defined \n",
    "by the dataframe object itself, which means that some data may not follow the expected rules. \n",
    "\n",
    "In `pandera` we can explicitly define these rules in schemas, which specify types for dataframe-like objects, \n",
    "and then use these schemas to assert properties about data at runtime and try parsing it into a desired state.\n",
    "\n",
    "Let's use a practical example. Suppose you're working with a transactions dataset of grocery `item`s and \n",
    "their associated `price`s. With these two categories we can make assumptions about the data and datatypes we \n",
    "expect in these fields. There may be a list of specific `item`s that are valid, or we can assume that any `price` \n",
    "should be greater than 0. We can state our assumptions about the data in `pandera` by writing a `Schema`, which can \n",
    "be defined in a `class`, as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9a4eef5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "import pandera.pandas as pa\n",
    "from pandera.typing import DataFrame, Series\n",
    "\n",
    "\n",
    "class Schema(pa.DataFrameModel):\n",
    "    item: Series[str] = pa.Field(isin=[\"apple\", \"orange\"], coerce=True)\n",
    "    price: Series[float] = pa.Field(gt=0, coerce=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "01cb7eb0",
   "metadata": {},
   "source": [
    "You can see that the `Schema` class inherits from [`pandera.DataFrameModel`](https://pandera.readthedocs.io/en/stable/reference/generated/pandera.api.pandas.model.DataFrameModel.html#pandera.api.pandas.model.DataFrameModel),\n",
    "and defines two fields: `item` and `price`. For each of these fields, `pandera` provides a \n",
    "flexible and concise way to specify the expected datatype: `str` for `item` and `float` for `price`. \n",
    "\n",
    "Other properties can also be set. For this example, we assumed that there might be a specific list of `item`s \n",
    "that are valid, or that a `price` should be greater than 0. These properties are defined in the `Schema`. \n",
    "In the code above, we use set equivalence for the `item` field with `isin=...` to specify valid options from a list, \n",
    "and we use value ranges for the `price` field with `gt=...` to specify a valid numeric range. These are only a couple \n",
    "examples of [property methods](https://pandera.readthedocs.io/en/stable/reference/generated/pandera.api.checks.Check.html#pandera.api.checks.Check) \n",
    "that can be asserted.\n",
    "\n",
    "Setting `coerce=True` will cause pandera to parse the columns into the expected datatypes, giving you the ability to \n",
    "ensure that data flowing through your pipeline is of the expected type."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e9efd661",
   "metadata": {},
   "source": [
    "## Runtime DataFrame Value Checks\n",
    "\n",
    "We can now use the `Schema` class to validate data passing through a function. In the example below, consider the \n",
    "function `add_sales_tax`, which will take the hypothetical grocery data and calculate the sales tax from the `price`, \n",
    "returning a new dataframe with the additional information in a new column.\n",
    "\n",
    "You can see why data validation would be important here. If the value in the `price` field is not the right datatype or \n",
    "is not greater than 0--as specified by the `Schema`--it will cause errors or further corrupt the data in any additional processes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa81fa0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pa.check_types(lazy=True)\n",
    "def add_sales_tax(data: DataFrame[Schema]):\n",
    "    # creates a new column in the data frame that calculates prices after sales tax\n",
    "    data[\"after_tax\"] = data[\"price\"] + (data[\"price\"] * 0.06)\n",
    "    return data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f5450ba9",
   "metadata": {},
   "source": [
    "As you will see when you run the code below, using the `@pa.check_types` [function decorator](https://pandera.readthedocs.io/en/stable/reference/decorators.html#decorators) and specifying the `data: DataFrame[Schema]` annotation in the function parameter \n",
    "will ensure that dataframe inputs are validated at runtime before being passed into the `add_sales_tax` function body.\n",
    "\n",
    "By providing the `lazy=True` option in the `check_types` decorator, we're telling `pandera` to validate all field properties \n",
    "before raising a `SchemaErrors` exception.\n",
    "\n",
    "With valid data, calling `add_sales_tax` shouldn't be a problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffd04ef5",
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_data = pd.DataFrame.from_records(\n",
    "    [{\"item\": \"apple\", \"price\": 0.5}, {\"item\": \"orange\", \"price\": 0.75}]\n",
    ")\n",
    "\n",
    "add_sales_tax(valid_data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "accb1038",
   "metadata": {},
   "source": [
    "With invalid data, however, `pandera` will raise a `SchemaErrors` exception:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1286e5e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "invalid_data = pd.DataFrame.from_records(\n",
    "    [{\"item\": \"applee\", \"price\": 0.5}, {\"item\": \"orange\", \"price\": -1000}]\n",
    ")\n",
    "\n",
    "try:\n",
    "    add_sales_tax(invalid_data)\n",
    "except pa.errors.SchemaErrors as exc:\n",
    "    display(exc.failure_cases)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a7595be4",
   "metadata": {},
   "source": [
    "The `exc.failure_cases` attribute in our `except` clause points to a dataframe that contains metadata about the failure \n",
    "cases that occurred when validating the data.\n",
    "\n",
    "We can see that row index `0` had a failure case in the misspelling of `applee` in the `item` column, which failed the \n",
    "`isin({\"apple\", \"orange\"})` check for that field.\n",
    "\n",
    "We can also see the row index `1` had a failure case of `-1000.0` in the `price` column, which failed the `gt=0` check for that field."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9e6e6099",
   "metadata": {},
   "source": [
    "## In-line Validation\n",
    "\n",
    "You can also use `Schema` classes to validate data in-line by calling the `validate` method, \n",
    "rather than at runtime as a part of a function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8d6ee48",
   "metadata": {},
   "outputs": [],
   "source": [
    "Schema.validate(valid_data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f35012b2",
   "metadata": {},
   "source": [
    "This gives you ultimate flexibility on where you want to validate data in your code."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "afc60bb7",
   "metadata": {},
   "source": [
    "## Schemas as Data Quality Checkpoints\n",
    "\n",
    "With `pandera`, you can use inheritance to indicate changes in the contents of a dataframe that some function \n",
    "has to implement. \n",
    "\n",
    "In the grocery example, let's assume we want to set an expiry date for each `item` in our list, but we want to \n",
    "validate the new data before AND after adding this new field, which means our schema will need to be different \n",
    "for the data in different points in the program. To accomplish this, first, we would build a second class that \n",
    "inherits from the original `Schema` class, as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c2a5bab",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Schema(pa.DataFrameModel):\n",
    "    item: Series[str] = pa.Field(isin=[\"apple\", \"orange\"], coerce=True)\n",
    "    price: Series[float] = pa.Field(gt=0, coerce=True)\n",
    "\n",
    "\n",
    "class TransformedSchema(Schema):\n",
    "    expiry: Series[pd.Timestamp] = pa.Field(coerce=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b0a80046",
   "metadata": {},
   "source": [
    "`TransformedSchema` will inherit the class attributes defined in `Schema`, with an additional `expiry` datetime field. \n",
    "In this case, we are asserting only a datatype of `Timestamp` on the `expiry` field.\n",
    "\n",
    "Now we can implement a function that performs the transformation needed to connect these two schemas. \n",
    "\n",
    "The `transform_data` function below takes a dataframe object and a list of `datetime`s and returns the input dataframe \n",
    "with a new column for `expiry` populated with the values of the `datetime` list argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f57e595e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "\n",
    "@pa.check_types(lazy=True)\n",
    "def transform_data(\n",
    "    data: DataFrame[Schema],\n",
    "    expiry: list[datetime],\n",
    ") -> DataFrame[TransformedSchema]:\n",
    "    return data.assign(expiry=expiry)\n",
    "\n",
    "\n",
    "transform_data(valid_data, [datetime.now()] * valid_data.shape[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a22d9462",
   "metadata": {},
   "source": [
    "Now every time we call the `transform_data` function, not only is the `data` input argument validated with the `Schema`, \n",
    "but the output dataframe is validated against `TransformedSchema`.\n",
    "\n",
    "In addition to catching value errors, this also allows you to catch bugs in your data transformation code more easily. \n",
    "Observe the buggy code below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9fa983e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pa.check_types(lazy=True)\n",
    "def transform_data(\n",
    "    data: DataFrame[Schema],\n",
    "    expiry: List[datetime],\n",
    ") -> DataFrame[TransformedSchema]:\n",
    "    return data.assign(expiryy=expiry)  # typo bug: 🐛\n",
    "\n",
    "\n",
    "try:\n",
    "    transform_data(valid_data, [datetime.now()] * valid_data.shape[0])\n",
    "except pa.errors.SchemaErrors as exc:\n",
    "    display(exc.failure_cases)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7fcd0187",
   "metadata": {},
   "source": [
    "The `failure_cases` dataframe is telling us in the `check` column that the core `column_in_dataframe` check is failing \n",
    "because the `expiry` column is not present in the output dataframe.\n",
    "\n",
    "Observe how the `schema_context` and `column` values in this `failure_cases` dataframe compare with those of the invalid \n",
    "data in the above examples. This shows the versitility of the error catching using `pandera`s `Schema`s. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "73b569b1",
   "metadata": {},
   "source": [
    "## Bonus: The Object-based API\n",
    "\n",
    "In the examples above, we've talked about dataframe schemas using the `DataFrameModel` or class-based API. \n",
    "However, `pandera` also provides an object-based API for defining dataframe schemas.\n",
    "\n",
    "While the [`DataFrameModel`](https://pandera.readthedocs.io/en/stable/dataframe_models.html) class-based API \n",
    "is closer in spirit to `dataclasses` and `pydantic`, which use Python classes to express complex data types , the\n",
    "object-based [`DataFrameSchema`](https://pandera.readthedocs.io/en/stable/dataframe_schemas.html) API enables \n",
    "you to transform your schema definition on the fly.\n",
    "\n",
    "Consider the difference between the class-based API and the equivalent object-based API syntax below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7597c976",
   "metadata": {},
   "outputs": [],
   "source": [
    "# class-based API\n",
    "class Schema(pa.DataFrameModel):\n",
    "    item: Series[str] = pa.Field(isin=[\"apple\", \"orange\"], coerce=True)\n",
    "    price: Series[float] = pa.Field(gt=0, coerce=True)\n",
    "\n",
    "\n",
    "# object-based API\n",
    "schema = pa.DataFrameSchema(\n",
    "    {\n",
    "        \"item\": pa.Column(\n",
    "            str, pa.Check.isin([\"apple\", \"orange\"]), coerce=True\n",
    "        ),\n",
    "        \"price\": pa.Column(float, pa.Check.gt(0), coerce=True),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23b06944",
   "metadata": {},
   "source": [
    "In the object-based API, you can add, remove, and update columns as you want, \n",
    "just as you would to a standard dataframe object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ef5399",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformed_schema = schema.add_columns({\"expiry\": pa.Column(pd.Timestamp)})\n",
    "schema.remove_columns([\"item\"])  # remove the \"item\" column\n",
    "schema.update_column(\n",
    "    \"price\", dtype=int\n",
    ")  # update the datatype of the \"price\" column to integer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ae1c08bb",
   "metadata": {},
   "source": [
    "You can use `DataFrameSchema`s to validate data just like `DataFrameModel` subclasses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80104780",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema.validate(valid_data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e913b9f9",
   "metadata": {},
   "source": [
    "And, similar to the `check_types` decorator, you can use the` check_io` decorator to validate inputs \n",
    "and outputs of your functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5af85e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pa.check_io(data=schema, out=transformed_schema)\n",
    "def fn(data, expiry):\n",
    "    return data.assign(expiry=expiry)\n",
    "\n",
    "\n",
    "fn(valid_data, [datetime.now()] * valid_data.shape[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bcecf627",
   "metadata": {},
   "source": [
    "### When to Use `DataFrameSchema` vs. `DataFrameModel`\n",
    "\n",
    "Practically speaking, the two ways of writing pandera schemas are completely equivalent, and using one over the \n",
    "other boils down to a few factors:\n",
    "\n",
    "1. Preference: some developers are more comfortable with one syntax over the other.\n",
    "2. The class-based API unlocks static type-checking of data via [mypy](https://pandera.readthedocs.io/en/stable/mypy_integration.html)\n",
    "   and integrates well with Python's type hinting system.\n",
    "3. The object-based API is good if you want to dynamically update your schema definition at runtime.\n",
    "\n",
    "At the end of the day, you can use them interchangeably in your applications."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9fe0f6aa",
   "metadata": {},
   "source": [
    "### What's Next?\n",
    "\n",
    "This notebook gave you a brief intro to Pandera, but this framework has a lot more to offer to help you test your data:\n",
    "\n",
    "- [Create in-line custom checks](https://pandera.readthedocs.io/en/stable/checks.html)\n",
    "- [Register custom checks](https://pandera.readthedocs.io/en/stable/extensions.html)\n",
    "- [Define statistical hypothesis tests](https://pandera.readthedocs.io/en/stable/hypothesis.html)\n",
    "- [Bootstrap schemas with data profiling](https://pandera.readthedocs.io/en/stable/schema_inference.html)\n",
    "- [Synthesize fake data for unit testing](https://pandera.readthedocs.io/en/stable/data_synthesis_strategies.html)\n",
    "- [Scale Validation with Distributed DataFrames](https://pandera.readthedocs.io/en/stable/supported_libraries.html#)\n",
    "- [Integrate with the Python Ecosystem](https://pandera.readthedocs.io/en/stable/integrations.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('pandera')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "vscode": {
   "interpreter": {
    "hash": "6bbffad058ea4c1a3a14a7a497e5fe0053e4f65d400b1604fe633506406cf190"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
